# coding=utf-8
__author__ = 'vivienfanghua'
from osgeo import ogr
from PIL import Image
from PIL import ImageDraw
from pyproj import Proj
import psycopg2
import traceback
class Map:
    '''
    这是一个绘制地图的类

    '''

    driver = ogr.GetDriverByName("ESRI Shapefile")
    conn=psycopg2.connect("dbname=gis user=postgres password=123456")
    p = Proj('+proj=aea +lon_0=105 +lat_1=25 +lat_2=47 +ellps=krass')
    cursor=conn.cursor()
    screenWidth=5000
    screenHeight=5000
    # left=73.557702
    # right=134.773925
    # bottom=15.780000
    # top=53.560861
    left=-3217402.170274192
    right=2233575.528979472
    bottom=1698930.774764923
    top=5055133.401208647
    im = Image.new("RGB", (screenWidth, screenHeight), "white")
    imDraw=ImageDraw.Draw(im)
    fileList=[]
    dataSourceList=[]
    layerList=[]
    def __init__(self):
        self.dataSourceList=[]
        self.layerList=[]
    def setup(self,attr):
        self.screenWidth=attr['width']
        self.screenHeight=attr['height']
        self.left=attr['left']
        self.right=attr['right']
        self.bottom=attr['bottom']
        self.up=attr['up']
    def readShp(self,*args):
        '''
        :param args: 任意多个shp文件名
        :return:无
        '''
        try:
            for i in args:
                dataSource=self.driver.Open(i,0)
                self.dataSourceList.append(dataSource)
                self.layerList.append(dataSource.GetLayer())
        except Exception as e:
            print(e.message)
    def shape2pgis(self,file,name):
        try:

            dataSource = self.driver.Open(file, 0)
            layer=dataSource.GetLayer()
            # self.cursor.execute("DROP TABLE IF EXISTS "+name)
            # self.cursor.execute("CREATE TABLE " + name)
            for i in range(layer.GetFeatureCount()):
                feature = layer.GetFeature(i)
                # name = feature.GetField("NAME")
                wkt = feature.GetGeometryRef().ExportToWkt()
                self.cursor.execute(
                    "INSERT INTO "+name+"(gamid, iso, pctbach, name_engli, name_iso, name_fao, name_local, name_obsol, name_varia) VALUES (" + '\'' +
                    feature.GetField[1] + '\'' + ", " + str(feature.GetField[2]) + ", " + str(feature.GetField[3]) + ", " + str(feature.GetField[4]) + ", " + str(
                        feature.GetField[5]) + ", " + str(feature.GetField[6]) + ", " + str(feature.GetField[7]) + ", " + str(feature.GetField[8]) +
                    ", ST_GeometryFromText(" + "'" + wkt + "', 4326))")
                print("INSERT "+str(i))
        except Exception as e:
            traceback.print_exc()
    def pgis2shape(self,name,file):
        try:
            out_ds = self.driver.CreateDataSource(file)
            out_layer = out_ds.CreateLayer("polygon", None, ogr.wkbPolygon)
            fd = ogr.FieldDefn('name', ogr.OFTString)
            out_layer.CreateField(fd)
            self.cursor.execute("SELECT name_iso,ST_AsText(geom) from public."+name+";")
            shp=self.cursor.fetchall()
            for row in shp:
                geo=ogr.CreateGeometryFromWkt(row[1])
                featDef = ogr.Feature(out_layer.GetLayerDefn())
                featDef.SetGeometry(geo)
                featDef.SetField('name', row[0])
                out_layer.CreateFeature(featDef)
            out_ds.Destroy()
                # featDef.SetField('name',)
        except Exception as e:
            traceback.print_exc()
    def World2Screen(self,x,y):
        x,y=self.p(x,y)
        worldSize=max(self.right-self.left,self.top-self.bottom)
        clientSize=min(self.screenHeight,self.screenWidth)
        scale=1.0*worldSize/clientSize
        return ((x-self.left)/scale,self.screenHeight-(y-self.bottom)/scale)
    def drawPolygon(self,geom):
        try:
            for i in range(geom.GetGeometryCount()):
                ring=geom.GetGeometryRef(i)
                lineString=ring.GetGeometryRef(0)
                pointCount=lineString.GetPointCount()
                coords=[]
                for p in xrange(pointCount):
                    lon, lat, z = lineString.GetPoint(p)
                    coords.append(self.World2Screen(lon,lat))
                self.imDraw.polygon(coords,fill=(255,0,0),outline=(50,50,50))
        except Exception as e:
            print(e.message)
    def drawPolyline(self,geom):
        pass
    def draw(self):
        for layer in self.layerList:
            for feature in layer:
                geom = feature.GetGeometryRef()
                type=geom.GetGeometryName()
                if type=="MULTIPOLYGON":
                    self.drawPolygon(geom)
                elif type=="MULTIPOLYLINE":
                    self.drawPolyline(geom)
        self.im.show()
    def drawPostGIS(self,name):
        try:
            self.cursor.execute("SELECT ST_AsText(geom) from public."+name+";")
            shp=self.cursor.fetchall()
            for row in shp:
                geo=ogr.CreateGeometryFromWkt(row[0])
                type=geo.GetGeometryName()
                if type=="MULTIPOLYGON":
                    self.drawPolygon(geo)
                elif type=="MULTIPOLYLINE":
                    self.drawPolyline(geo)
            self.im.show()
        except Exception as e:
            print(e.message)
    def export(self,file):
        try:
            self.im.save(file)
        except Exception as e:
            print(e.message)


def main():
    map=Map()
    # map.readShp('D:/CHN_adm/CHN_adm2.shp')
    # map.draw()
    map.drawPostGIS("chn_adm2")
    # map.shape2pgis("/home/vivien/PycharmProjects/python-learn/CHN_adm/CHN_adm0.shp","CHN_adm0")
    # map.pgis2shape("CHN_adm0","/home/vivien/PycharmProjects/python-learn/CHN_adm/test.shp")
    map.export("/home/vivien/result.jpg")

if __name__=="__main__":
    main()